#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：PythonProjects
@File    ：predict.py
@IDE     ：PyCharm
@Author  ：pipibao
@Date    ：2021/7/5 下午9:54
'''
import tensorflow.compat.v1 as tf

tf.disable_v2_behavior()
import numpy as np
import cv2 as cv
import sys
import os
import random

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 设置警告等级


# 定义卷积函数
def conv_layer(inputs, W, b, conv_strides, kernel_size, pool_strides, padding):
    L1_conv = tf.nn.conv2d(inputs, W, strides=conv_strides, padding=padding)  # 卷积操作
    L1_relu = tf.nn.relu(L1_conv + b)  # 激活函数RELU
    return tf.nn.max_pool(L1_relu, ksize=kernel_size, strides=pool_strides, padding='SAME')


# 定义全连接函数
def full_connect(inputs, W, b):
    return tf.nn.relu(tf.matmul(inputs, W) + b)


# 定义第一个预测函数
def predicts(Path="test\\"):
    PROVINCES = (
        "川", "鄂", "赣", "甘", "贵", "桂", "黑", "沪", "冀", "津", "京", "吉", "辽", "鲁", "蒙", "闽", "宁", "青", "琼", "陕", "苏", "晋",
        "皖",
        "湘", "新", "豫", "渝", "粤", "云", "藏", "浙")
    nProvinceIndex = 0
    SAVER_DIR = "train_saver\\chinese\\"
    # 新建一个图
    g1 = tf.Graph()
    with g1.as_default():
        x = tf.placeholder(tf.float32, shape=[None, 1024])  # None表示batch size的大小，这里可以是任何数，因为不知道待训练的图片数，SIZE指图片的大小
        y_ = tf.placeholder(tf.float32, shape=[None, 31])  # 输出标签的占位
        x_image = tf.reshape(x, [-1, 32, 32, 1])  # 生成一个四维的数组
        sess1 = tf.Session(graph=g1)
        saver = tf.train.import_meta_graph("%smodel.ckpt.meta" % (SAVER_DIR))
        # model_file = "%smodel.ckpt"%(SAVER_DIR)
        model_file = tf.train.latest_checkpoint(SAVER_DIR)  # 找出所有模型中最新的模型
        saver.restore(sess1, model_file)  # 恢复模型，相当于加载模型
        # 第一个卷积层
        W_conv1 = sess1.graph.get_tensor_by_name("W_conv1:0")
        b_conv1 = sess1.graph.get_tensor_by_name("b_conv1:0")
        conv_strides = [1, 1, 1, 1]
        kernel_size = [1, 2, 2, 1]
        pool_strides = [1, 2, 2, 1]
        L1_pool = conv_layer(x_image, W_conv1, b_conv1, conv_strides, kernel_size, pool_strides, padding='SAME')
        print("第一个卷积层")
        # 第二个卷积层
        W_conv2 = sess1.graph.get_tensor_by_name("W_conv2:0")
        b_conv2 = sess1.graph.get_tensor_by_name("b_conv2:0")
        conv_strides = [1, 1, 1, 1]
        kernel_size = [1, 2, 2, 1]
        pool_strides = [1, 2, 2, 1]
        L2_pool = conv_layer(L1_pool, W_conv2, b_conv2, conv_strides, kernel_size, pool_strides, padding='SAME')

        # 全连接层
        W_fc1 = sess1.graph.get_tensor_by_name("W_fc1:0")
        b_fc1 = sess1.graph.get_tensor_by_name("b_fc1:0")
        h_pool2_flat = tf.reshape(L2_pool, [-1, 8 * 8 * 24])
        h_fc1 = full_connect(h_pool2_flat, W_fc1, b_fc1)

        # dropout
        keep_prob = tf.placeholder(tf.float32)
        h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

        # readout层
        W_fc2 = sess1.graph.get_tensor_by_name("W_fc2:0")
        b_fc2 = sess1.graph.get_tensor_by_name("b_fc2:0")

        # 定义优化器和训练op
        conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

        for n in range(1, 2):
            path = Path + "%s.jpg" % (n)  # 测试图的路径
            img = cv.imread(path, 0)
            # cv.imshow('threshold',img)
            # cv.waitKey(0)
            height = img.shape[0]  # 行数
            width = img.shape[1]  # 列数

            img_data = [[0] * 1024 for i in range(1)]  # 创建一个数组，用于将输入的图片转换成数组形式
            for h in range(0, height):
                for w in range(0, width):
                    m = img[h][w]
                    if m > 150:
                        img_data[0][w + h * width] = 1
                    else:
                        img_data[0][w + h * width] = 0

            result = sess1.run(conv, feed_dict={x: np.array(img_data), keep_prob: 1.0})

            # 用于输出概率最大的3类
            max1 = 0
            max2 = 0
            max3 = 0
            max1_index = 0
            max2_index = 0
            max3_index = 0
            for j in range(31):
                if result[0][j] > max1:
                    max1 = result[0][j]
                    max1_index = j
                    continue
                if (result[0][j] > max2) and (result[0][j] <= max1):
                    max2 = result[0][j]
                    max2_index = j
                    continue
                if (result[0][j] > max3) and (result[0][j] <= max2):
                    max3 = result[0][j]
                    max3_index = j
                    continue
            nProvinceIndex = max1_index  # 最大概率的类
            print("概率：[%s %0.2f%%]    [%s %0.2f%%]    [%s %0.2f%%]" % (
                PROVINCES[max1_index], max1 * 100, PROVINCES[max2_index], max2 * 100, PROVINCES[max3_index],
                max3 * 100))

        print("省份简称是：%s" % PROVINCES[nProvinceIndex])
        return PROVINCES[nProvinceIndex], nProvinceIndex
        sess1.close()


# 定义第二个预测函数
def predictn(Path="test\\"):
    LETTERS_DIGITS = (
        "A", "B", "C", "D", "E", "F", "G", "H", "J", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "U", "V", "W", "X",
        "Y",
        "Z", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9")
    license_num = ""
    SAVER_DIR = "train_saver\\char\\"

    print("进入调用")
    g2 = tf.Graph()
    with g2.as_default():
        x = tf.placeholder(tf.float32, shape=[None, 1024])  # None表示batch size的大小，这里可以是任何数，因为不知道待训练的图片数，SIZE指图片的大小
        y_ = tf.placeholder(tf.float32, shape=[None, 34])  # 输出标签的占位
        x_image = tf.reshape(x, [-1, 32, 32, 1])  # 生成一个四维的数组
        sess2 = tf.Session(graph=g2)
        saver = tf.train.import_meta_graph("%smodel.ckpt.meta" % (SAVER_DIR))
        model_file = tf.train.latest_checkpoint(SAVER_DIR)  # 找出所有模型中最新的模型
        saver.restore(sess2, model_file)
        # 第一个卷积层
        W_conv1 = sess2.graph.get_tensor_by_name("W_conv1:0")
        b_conv1 = sess2.graph.get_tensor_by_name("b_conv1:0")
        conv_strides = [1, 1, 1, 1]
        kernel_size = [1, 2, 2, 1]
        pool_strides = [1, 2, 2, 1]
        L1_pool = conv_layer(x_image, W_conv1, b_conv1, conv_strides, kernel_size, pool_strides, padding='SAME')
        # 第二个卷积层
        W_conv2 = sess2.graph.get_tensor_by_name("W_conv2:0")
        b_conv2 = sess2.graph.get_tensor_by_name("b_conv2:0")
        conv_strides = [1, 1, 1, 1]
        kernel_size = [1, 2, 2, 1]
        pool_strides = [1, 2, 2, 1]
        L2_pool = conv_layer(L1_pool, W_conv2, b_conv2, conv_strides, kernel_size, pool_strides, padding='SAME')
        # 全连接层
        W_fc1 = sess2.graph.get_tensor_by_name("W_fc1:0")
        b_fc1 = sess2.graph.get_tensor_by_name("b_fc1:0")
        h_pool2_flat = tf.reshape(L2_pool, [-1, 8 * 8 * 24])
        h_fc1 = full_connect(h_pool2_flat, W_fc1, b_fc1)
        # dropout
        keep_prob = tf.placeholder(tf.float32)
        h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
        # readout层
        W_fc2 = sess2.graph.get_tensor_by_name("W_fc2:0")
        b_fc2 = sess2.graph.get_tensor_by_name("b_fc2:0")
        # 定义优化器和训练op
        conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)
        # 想尝试将城市代码和车牌后五位一起识别，因此可以将3-8改为2-8
        for n in range(2, 8):
            path = Path + "%s.jpg" % (n)
            img = cv.imread(path, 0)
            height = img.shape[0]
            width = img.shape[1]
            img_data = [[0] * 1024 for i in range(1)]
            for h in range(0, height):
                for w in range(0, width):
                    m = img[h][w]
                    if m > 150:
                        img_data[0][w + h * width] = 1
                    else:
                        img_data[0][w + h * width] = 0

            result = sess2.run(conv, feed_dict={x: np.array(img_data), keep_prob: 1.0})
            max1 = 0
            max2 = 0
            max3 = 0
            max1_index = 0
            max2_index = 0
            max3_index = 0
            for j in range(34):
                if result[0][j] > max1:
                    max1 = result[0][j]
                    max1_index = j
                    continue
                if (result[0][j] > max2) and (result[0][j] <= max1):
                    max2 = result[0][j]
                    max2_index = j
                    continue
                if (result[0][j] > max3) and (result[0][j] <= max2):
                    max3 = result[0][j]
                    max3_index = j
                    continue

            license_num = license_num + LETTERS_DIGITS[max1_index]
            print("概率：[%s %0.2f%%]    [%s %0.2f%%]    [%s %0.2f%%]" % (
                LETTERS_DIGITS[max1_index], max1 * 100, LETTERS_DIGITS[max2_index], max2 * 100,
                LETTERS_DIGITS[max3_index],
                max3 * 100))

        print("车牌编号是：%s" % license_num)
        return license_num
        sess2.close()


def predict(Path="test\\"):
    a, b = predicts(Path)
    c = predictn(Path)
    return str("车牌号为：" + a + c[0] + "·" + c[1:6])


if __name__ == "__main__":
    ans = predict()
    print(ans)
